"""
Symmetry breaking experiment for Allen-Cahn equation using DLR-FEM.

This script implements the kiss-bubble experiment from Section 4 of the paper,
which demonstrates the ability of DLR-FEM to capture symmetry breaking dynamics
in gradient flows with low-rank structure.
"""

import numpy as np
import matplotlib.pyplot as plt
from typing import Dict, List, Optional, Tuple
import time
import os
import sys

from ..main import DLRFEMSolver, DLRFEMConfig


class SymmetryBreakingExperiment:
    """
    Experiment class for studying symmetry breaking in Allen-Cahn equation.
    
    Implements the kiss-bubble initial condition and analyzes the dynamics
    of bubble coalescence and symmetry breaking.
    """
    
    def __init__(self, config: Optional[DLRFEMConfig] = None):
        """
        Initialize symmetry breaking experiment.
        
        Args:
            config: DLR-FEM configuration
        """
        if config is None:
            # Default configuration for symmetry breaking experiment
            self.config = DLRFEMConfig(
                domain_x=(0.0, 1.0),
                domain_y=(0.0, 1.0),
                nx=64,
                ny=64,
                poly_order=1,
                epsilon=0.02,
                T_final=1.0,
                dt=0.001,
                initial_rank=15,
                max_rank=50,
                truncation_tol=1e-12,
                adaptive=False,
                save_frequency=10,
                plot_frequency=50,
                verbose=True
            )
        else:
            self.config = config
        
        self.solver = DLRFEMSolver(self.config)
        
    def run_standard_experiment(self) -> Dict:
        """
        Run the standard kiss-bubble experiment from the paper.
        
        Uses the configuration from Eq. (4.2) with two bubbles positioned
        to eventually interact and break symmetry.
        
        Returns:
            Results dictionary
        """
        print("Running standard kiss-bubble symmetry breaking experiment...")
        
        # Standard kiss-bubble configuration
        bubble_centers = [(0.3, 0.5), (0.7, 0.5)]  # Two bubbles on horizontal line
        bubble_radius = 0.15
        
        results = self.solver.solve(
            ic_type="kiss_bubble",
            centers=bubble_centers,
            radius=bubble_radius,
            epsilon=self.config.epsilon
        )
        
        # Analyze symmetry breaking
        self._analyze_symmetry_breaking(results)
        
        return results
    
    def run_parameter_study(self, epsilon_values: List[float]) -> Dict:
        """
        Run parameter study varying interface thickness parameter ε.
        
        Args:
            epsilon_values: List of epsilon values to test
            
        Returns:
            Dictionary with results for each epsilon
        """
        print(f"Running parameter study for ε values: {epsilon_values}")
        
        all_results = {}
        
        for eps in epsilon_values:
            print(f"\n--- Running simulation with ε = {eps} ---")
            
            # Update configuration
            config = self.config
            config.epsilon = eps
            
            # Create new solver with updated epsilon
            solver = DLRFEMSolver(config)
            
            # Run simulation
            results = solver.solve(
                ic_type="kiss_bubble",
                centers=[(0.3, 0.5), (0.7, 0.5)],
                radius=0.15,
                epsilon=eps
            )
            
            all_results[eps] = results
        
        # Compare results
        self._compare_parameter_results(all_results)
        
        return all_results
    
    def run_bubble_configuration_study(self) -> Dict:
        """
        Study different bubble configurations and their evolution.
        
        Returns:
            Dictionary with results for different configurations
        """
        print("Running bubble configuration study...")
        
        configurations = {
            'symmetric': [(0.25, 0.5), (0.75, 0.5)],  # Symmetric on x-axis
            'offset': [(0.3, 0.4), (0.7, 0.6)],       # Slightly offset
            'close': [(0.4, 0.5), (0.6, 0.5)],        # Closer together
            'far': [(0.2, 0.5), (0.8, 0.5)],          # Farther apart
            'three_bubbles': [(0.25, 0.5), (0.5, 0.3), (0.75, 0.7)]  # Three bubbles
        }
        
        all_results = {}
        
        for name, centers in configurations.items():
            print(f"\n--- Running configuration: {name} ---")
            
            results = self.solver.solve(
                ic_type="kiss_bubble",
                centers=centers,
                radius=0.12 if name != 'three_bubbles' else 0.1,
                epsilon=self.config.epsilon
            )
            
            all_results[name] = results
        
        # Plot comparison
        self._plot_configuration_comparison(all_results)
        
        return all_results
    
    def _analyze_symmetry_breaking(self, results: Dict):
        """
        Analyze the symmetry breaking behavior.
        
        Args:
            results: Simulation results
        """
        print("\nAnalyzing symmetry breaking...")
        
        time_history = results['time_history']
        solution_history = results['solution_history']
        
        if len(solution_history) < 2:
            print("Not enough solution history for analysis")
            return
        
        # Compute symmetry measure over time
        symmetry_measures = []
        
        for sol in solution_history:
            # Measure left-right symmetry by comparing u(x,y) with u(1-x,y)
            sol_flipped = np.fliplr(sol)
            symmetry_error = np.linalg.norm(sol - sol_flipped, 'fro')
            symmetry_measures.append(symmetry_error)
        
        # Plot symmetry evolution
        plt.figure(figsize=(12, 8))
        
        # Symmetry measure over time
        plt.subplot(2, 2, 1)
        time_indices = np.arange(len(symmetry_measures)) * self.config.save_frequency
        actual_times = [time_history[min(i, len(time_history)-1)] for i in time_indices]
        plt.plot(actual_times, symmetry_measures, 'b-', linewidth=2)
        plt.xlabel('Time')
        plt.ylabel('Symmetry Error')
        plt.title('Symmetry Breaking Evolution')
        plt.grid(True)
        
        # Rank evolution
        plt.subplot(2, 2, 2)
        plt.plot(time_history, results['rank_history'], 'r-', linewidth=2)
        plt.xlabel('Time')
        plt.ylabel('Rank')
        plt.title('Rank Evolution')
        plt.grid(True)
        
        # Initial and final solutions
        plt.subplot(2, 2, 3)
        X, Y = results['mesh'].get_coordinates()
        plt.contourf(X, Y, solution_history[0], levels=20, cmap='coolwarm')
        plt.title('Initial Condition')
        plt.xlabel('x')
        plt.ylabel('y')
        plt.colorbar()
        plt.axis('equal')
        
        plt.subplot(2, 2, 4)
        plt.contourf(X, Y, solution_history[-1], levels=20, cmap='coolwarm')
        plt.title('Final Solution')
        plt.xlabel('x')
        plt.ylabel('y')
        plt.colorbar()
        plt.axis('equal')
        
        plt.tight_layout()
        plt.show()
        
        # Print analysis summary
        initial_symmetry = symmetry_measures[0]
        final_symmetry = symmetry_measures[-1]
        max_symmetry_break = max(symmetry_measures)
        
        print(f"Symmetry Analysis Results:")
        print(f"  Initial symmetry error: {initial_symmetry:.2e}")
        print(f"  Final symmetry error: {final_symmetry:.2e}")
        print(f"  Maximum symmetry break: {max_symmetry_break:.2e}")
        print(f"  Symmetry breaking factor: {final_symmetry / initial_symmetry:.1f}")
        
        # Detect time of symmetry breaking
        threshold = initial_symmetry * 2  # When error doubles
        break_idx = np.where(np.array(symmetry_measures) > threshold)[0]
        if len(break_idx) > 0:
            break_time = actual_times[break_idx[0]]
            print(f"  Symmetry breaking detected at t ≈ {break_time:.3f}")
    
    def _compare_parameter_results(self, all_results: Dict):
        """
        Compare results across different parameter values.
        
        Args:
            all_results: Dictionary of results for different parameters
        """
        print("\nComparing parameter study results...")
        
        fig, axes = plt.subplots(2, 3, figsize=(18, 12))
        epsilon_values = list(all_results.keys())
        
        # Plot final solutions for different epsilon values
        for i, eps in enumerate(epsilon_values[:6]):  # Show up to 6 values
            if i >= 6:
                break
                
            row = i // 3
            col = i % 3
            
            results = all_results[eps]
            final_solution = results['solution']
            X, Y = results['mesh'].get_coordinates()
            
            ax = axes[row, col]
            im = ax.contourf(X, Y, final_solution, levels=20, cmap='coolwarm')
            ax.set_title(f'ε = {eps}')
            ax.set_xlabel('x')
            ax.set_ylabel('y')
            ax.axis('equal')
            plt.colorbar(im, ax=ax)
        
        # Hide unused subplots
        for i in range(len(epsilon_values), 6):
            row = i // 3
            col = i % 3
            axes[row, col].axis('off')
        
        plt.tight_layout()
        plt.show()
        
        # Compare rank evolution
        plt.figure(figsize=(12, 5))
        
        plt.subplot(1, 2, 1)
        for eps in epsilon_values:
            results = all_results[eps]
            plt.plot(results['time_history'], results['rank_history'], 
                    label=f'ε = {eps}', linewidth=2)
        plt.xlabel('Time')
        plt.ylabel('Rank')
        plt.title('Rank Evolution Comparison')
        plt.legend()
        plt.grid(True)
        
        # Compare final ranks
        plt.subplot(1, 2, 2)
        final_ranks = [all_results[eps]['rank_history'][-1] for eps in epsilon_values]
        plt.plot(epsilon_values, final_ranks, 'bo-', linewidth=2, markersize=8)
        plt.xlabel('ε')
        plt.ylabel('Final Rank')
        plt.title('Final Rank vs Interface Parameter')
        plt.grid(True)
        
        plt.tight_layout()
        plt.show()
    
    def _plot_configuration_comparison(self, all_results: Dict):
        """
        Plot comparison of different bubble configurations.
        
        Args:
            all_results: Results for different configurations
        """
        print("\nPlotting configuration comparison...")
        
        n_configs = len(all_results)
        fig, axes = plt.subplots(2, n_configs, figsize=(4*n_configs, 8))
        
        if n_configs == 1:
            axes = axes.reshape(2, 1)
        
        for i, (name, results) in enumerate(all_results.items()):
            X, Y = results['mesh'].get_coordinates()
            
            # Initial condition
            if len(results['solution_history']) > 0:
                ax = axes[0, i]
                im = ax.contourf(X, Y, results['solution_history'][0], 
                               levels=20, cmap='coolwarm')
                ax.set_title(f'{name} - Initial')
                ax.set_xlabel('x')
                ax.set_ylabel('y')
                ax.axis('equal')
                plt.colorbar(im, ax=ax, shrink=0.8)
            
            # Final solution
            ax = axes[1, i]
            im = ax.contourf(X, Y, results['solution'], 
                           levels=20, cmap='coolwarm')
            ax.set_title(f'{name} - Final')
            ax.set_xlabel('x')
            ax.set_ylabel('y')
            ax.axis('equal')
            plt.colorbar(im, ax=ax, shrink=0.8)
        
        plt.tight_layout()
        plt.show()
        
        # Rank comparison
        plt.figure(figsize=(10, 6))
        for name, results in all_results.items():
            plt.plot(results['time_history'], results['rank_history'], 
                    label=name, linewidth=2)
        plt.xlabel('Time')
        plt.ylabel('Rank')
        plt.title('Rank Evolution - Configuration Comparison')
        plt.legend()
        plt.grid(True)
        plt.show()
    
    def generate_animation(self, results: Dict, filename: str = "symmetry_breaking.gif"):
        """
        Generate animation of the symmetry breaking process.
        
        Args:
            results: Simulation results
            filename: Output filename for animation
        """
        print(f"Generating animation: {filename}")
        
        from matplotlib.animation import FuncAnimation
        
        X, Y = results['mesh'].get_coordinates()
        solution_history = results['solution_history']
        time_history = results['time_history']
        
        if len(solution_history) < 2:
            print("Not enough solution history for animation")
            return
        
        fig, ax = plt.subplots(figsize=(8, 6))
        
        # Set up the plot
        vmin = min(np.min(sol) for sol in solution_history)
        vmax = max(np.max(sol) for sol in solution_history)
        
        def animate(frame):
            ax.clear()
            
            # Get time index
            time_idx = frame * self.config.save_frequency
            actual_time = time_history[min(time_idx, len(time_history)-1)]
            
            # Plot solution
            im = ax.contourf(X, Y, solution_history[frame], 
                           levels=20, cmap='coolwarm', vmin=vmin, vmax=vmax)
            ax.set_title(f'DLR-FEM Symmetry Breaking: t = {actual_time:.3f}')
            ax.set_xlabel('x')
            ax.set_ylabel('y')
            ax.axis('equal')
            
            return im.collections
        
        # Create animation
        anim = FuncAnimation(fig, animate, frames=len(solution_history), 
                           interval=200, blit=False, repeat=True)
        
        # Save animation
        try:
            anim.save(filename, writer='pillow', fps=5)
            print(f"Animation saved as {filename}")
        except Exception as e:
            print(f"Could not save animation: {e}")
            print("Showing animation instead...")
            plt.show()


def main():
    """Run the complete symmetry breaking experiment suite."""
    print("DLR-FEM Symmetry Breaking Experiment Suite")
    print("=" * 50)
    
    # Create experiment with default configuration
    experiment = SymmetryBreakingExperiment()
    
    # 1. Standard experiment
    print("\n1. Running standard kiss-bubble experiment...")
    standard_results = experiment.run_standard_experiment()
    
    # 2. Parameter study
    print("\n2. Running parameter study...")
    epsilon_values = [0.01, 0.02, 0.03, 0.04]
    param_results = experiment.run_parameter_study(epsilon_values)
    
    # 3. Configuration study
    print("\n3. Running bubble configuration study...")
    config_results = experiment.run_bubble_configuration_study()
    
    # 4. Generate animation for standard experiment
    print("\n4. Generating animation...")
    experiment.generate_animation(standard_results)
    
    print("\nSymmetry breaking experiment suite completed!")
    
    return {
        'standard': standard_results,
        'parameter_study': param_results,
        'configuration_study': config_results
    }


if __name__ == "__main__":
    results = main()